library(tidyverse)
library(hierfstat)
library(ggpubr)

g2r_position <- 106204113

# ensembl data sliced 1000 genomes vcf
chr14 <- read_delim('Archive/covid19/data/14.106070000-106260000.ALL.chr14.phase3_shapeit2_mvncall_integrated_v5b.20130502.genotypes.vcf.gz',
               comment = '##')

sub_meta <- read.delim('Archive/covid19/data/integrated_call_samples_v3.20130502.ALL.panel.txt') |>
  select(1:3)

non_cdx_meta <- sub_meta |>
  filter(pop == 'CDX') |>
  pull(sample)

chr14 |>
  select(-any_of(cdx_meta)) |>
  write_delim('Archive/covid19/results/1kgp.cdx.ighg1.gr.vcf')

filter_mask <- read_delim('Archive/covid19/data/chr14_mask.bed.gz', col_names = FALSE) |>
  filter(X2 > 10607e4 & X3 < 10626e4)

# calc Fst --------
fstat <- chr14 |>
  distinct(POS, .keep_all = TRUE) |>
  select(c(POS, 10:2513)) |>
  mutate(across(everything(), \(x)str_remove(x, '\\|'))) |>
  column_to_rownames('POS') |>
  t() |>
  as.data.frame() |>
  rownames_to_column('sample') |>
  left_join(sub_meta[1:2])

fstat_diploid <- fstat |>
  mutate(pop_id = case_when(
    pop == 'CDX' ~ 1,
    !str_detect(pop, 'CHS|CHB|KHV|JPT') ~ 2,
    .default = 0
  )) |>
  filter(pop_id != 0) |>
  column_to_rownames('sample') |>
  select(-pop) |>
  relocate(pop_id) |>
  mutate(across(everything(), as.numeric))

fst_res <- basic.stats(fstat_diploid)

CDX_other_fst <- fst_res$perloc |>
  select(Fst) |>
  rownames_to_column('position') |>
  as_tibble() |>
  mutate(position = as.numeric(position)) |>
  filter(!is.na(Fst))

write_csv(CDX_other_fst, 'Archive/covid19/results/CDX_vs_non-EAS_Fst.csv')

CDX_other_fst <- read_csv('Archive/covid19/results/CDX_vs_non-EAS_Fst.csv')

final_block <- read_csv('Archive/covid19/results/final_block_candidate13.csv')

highest_CDX_fst <- CDX_other_fst |>
  slice_max(Fst, n = 20)

block_low_bound <- 106192543
block_high_bound <- 106253349

# plot Fst points ----------
CDX_other_fst |>
  filter(position > 10607e4) |>
  ggplot(aes(position, Fst)) +
  geom_point() +
  geom_point(data = quality_fst, color = 'blue', size = 3) +
  geom_point(data = filter(CDX_other_fst, position == g2r_position), color = 'red', size = 3) +
  geom_point(data = filter(CDX_other_fst, position == g2r_position), color = 'red', size = 5, shape = "circle open") +
  theme_pubr() +
  labs_pubr() +
  geom_vline(xintercept = c(106202680, 106209408), color = 'orange', linetype = 'dashed') +
  ylab('Fst (CDX vs non-EAS)')

# find high quality sites on Vindija chr14 -----------
find_high_quality <- function(val){
  filter_mask |>
    add_column(y = val) |>
    pull(in_range = (y - X2) * (y - X3) <= 0) |>
    some(isTRUE)
}

find_high_quality(106241806)

mask_res <- CDX_other_fst$position |>
  map_lgl(find_high_quality, .progress = TRUE)

quality_fst <- CDX_other_fst |>
  mutate(high_quality = mask_res) |>
  filter(high_quality)

quality_high_fst <- slice_max(quality_fst, Fst, n = 10)

write_csv(quality_fst, 'Archive/covid19/results/CDX_quality_Fst.csv')

quality_fst |>
  ggplot(aes(position, Fst)) +
  geom_point() +
  geom_point(data = quality_high_fst, color = 'red', size = 3) +
  theme_pubr() +
  labs_pubr() +
  geom_vline(xintercept = c(106202680, 106209408), color = 'orange', linetype = 'dashed') +
  ylab('Fst (CDX vs non-EAS)')

# sample in position window to plot cleaner
CDX_other_fst |>
  filter(!is.na(Fst)) |>
  group_by(position_tile = ntile(position, 100)) |>
  reframe(median_pos = median(position),
          sig_fst = max(Fst)) |>
  ggplot(aes(median_pos, sig_fst)) +
  geom_line()

# sub population --------
allele_sub <- chr14 |>
  select(-c(INFO, ID, QUAL, FILTER, FORMAT)) |>
  pivot_longer(cols = 5:2508, names_to = 'sample', values_to = 'allele') |>
  mutate(allele = case_match(allele,
                             '1|1' ~ 1,
                             c('1|0', '0|1') ~ 0.5,
                             '0|0' ~ 0)) |>
  filter(!is.na(allele)) |>
  left_join(sub_meta)

allele_sub |>
  filter(pop == 'CDX') |>
  group_by(POS) |>
  summarize(x = sum(allele) * 2) |>
  filter(x != 0) |>
  mutate(n = 186, folded = 0) |>
  rename(position = POS) |>
  write_tsv('Archive/covid19/results/cdx-ighg1-sweep-finder.txt')

# re-draw G2R 1kgp freq plot in MS -------
swb_color <- c('grey','blue','red')

allele_sub |>
  filter(POS == g2r_position) |>
  mutate(allele = case_match(allele,
                             0 ~ 'CC',
                             0.5 ~ 'CT',
                             1 ~'TT'),
         super_pop = fct_relevel(super_pop, 'EAS', 'SAS')) |>
  ggplot(aes(super_pop, fill = allele)) +
  geom_bar(position = 'fill') +
  labs_pubr() +
  theme_pubr() +
  labs(x = 'Superpopulation', y = 'Frequency') +
  scale_fill_manual(values = swb_color)

g1 <- last_plot()

allele_sub |>
  filter(POS == g2r_position) |>
  mutate(allele = case_match(allele,
                             0 ~ 'CC',
                             0.5 ~ 'CT',
                             1 ~'TT'),
         pop = fct_relevel(pop, 'CDX', 'KHV', 'CHS', 'CHB', 'BEB', 'JPT')) |>
  ggplot(aes(pop, fill = allele)) +
  geom_bar(position = 'fill') +
  labs_pubr() +
  theme_pubr() +
  labs(x = 'Population', y = 'Frequency') +
  scale_fill_manual(values = swb_color)

g2 <- last_plot()

g1 + g2 +
  patchwork::plot_layout(widths = c(1,2), guides = 'collect')

ggsave('Archive/covid19/figures/FigS5_g2r_1kgp_pop.pdf',width = 20, height = 7)

global_maf <- allele_sub |>
  group_by(POS) |>
  summarise(MAF = mean(allele))

# quantify CDX's MAF divergence
diverge_sub <- allele_sub |>
  filter(pop == 'CDX') |>
  group_by(POS) |>
  summarise(CDX_MAF = mean(allele)) |>
  left_join(global_maf) |>
  mutate(CDX_diverge = abs(CDX_MAF - MAF))

set.seed(42)

sites_as_ctrl <- diverge_sub |>
  filter((POS < block_low_bound | POS > block_high_bound) &
           CDX_MAF > 0.01 & MAF > 0.01) |>
  slice_sample(n = 30)

write_csv(sites_as_ctrl, 'Archive/covid19/results/CDX_LD_ctrl_sites.csv')

close_sub10 <- diverge_sub |>
  filter(POS > g2r_position - 15e3 & POS < g2r_position + 15e3) |>
  slice_max(CDX_diverge, n = 10)

most_sub20 <- diverge_sub |>
  slice_max(CDX_diverge, n = 20)

sub_color = c(colorRampPalette(c('red','orange'))(6),
              colorRampPalette(c('grey60','black'))(20))

maf_sub <- allele_sub |>
  group_by(POS, pop) |>
  summarise(MAF = mean(allele))

write_csv(maf_sub, 'Archive/covid19/results/1KGP_subpop_MAF.csv')

maf_sub |>
  filter(POS %in% quality_position) |>
  mutate(pop = fct_relevel(pop, 'CDX', 'KHV', 'CHS', 'CHB', 'BEB', 'JPT')) |>
  ggplot(aes(pop, MAF, fill = pop)) +
  geom_col() +
  facet_wrap(~POS) +
  scale_fill_manual(values = sub_color) +
  theme(axis.text.x = element_text(angle = 90, hjust = 1, vjust = .5)) +
  scale_x_discrete(breaks = NULL) +
  labs(fill = 'population')

allele_sub |>
  filter(POS %in% close_sub10$POS) |>
  group_by(POS, pop) |>
  summarise(MAF = mean(allele)) |>
  mutate(pop = fct_relevel(pop, 'CDX', 'KHV', 'CHS', 'CHB', 'BEB', 'JPT'),
         POS = as.character(POS) |> fct_reorder(MAF)) |>
  ggplot(aes(pop, MAF, fill = pop)) +
  geom_col() +
  facet_wrap(~POS) +
  scale_fill_manual(values = sub_color) +
  theme(axis.text.x = element_text(angle = 90, hjust = 1, vjust = .5)) +
  scale_x_discrete(breaks = NULL) +
  labs(fill = 'population')

new_colocate_pos <- c(106241806,
                      106253183,
                      106080190)

allele_sub |>
  filter(POS %in% new_colocate_pos) |>
  group_by(POS, pop) |>
  summarise(MAF = mean(allele)) |>
  mutate(pop = fct_relevel(pop, 'CDX', 'KHV', 'CHS', 'CHB', 'BEB', 'JPT'),
         POS = as.character(POS) |> fct_reorder(MAF)) |>
  ggplot(aes(pop, MAF, fill = pop)) +
  geom_col() +
  facet_wrap(~POS) +
  scale_fill_manual(values = sub_color) +
  theme(axis.text.x = element_text(angle = 90, hjust = 1, vjust = .5)) +
  scale_x_discrete(breaks = NULL) +
  labs(fill = 'population')

# visuallize position on genome
diverge_sub |>
  ggplot(aes(x = POS, y = CDX_diverge)) +
  geom_point() +
  geom_point(data = most_sub20, color = 'blue', size = 2) +
  geom_point(aes(106204113, 0.823), color = 'red', size = 3) +
  geom_vline(xintercept = c(106204113 + 15000, 106204113 - 15000), color = 'orange', linetype = 'dashed') +
  theme_pubr() +
  geom_rect(aes(xmin = 106202680, xmax = 106209408, ymin = -0.1, ymax = -0.05))

diverge_sub |>
  ggplot(aes(x = POS, y = CDX_diverge)) +
  geom_point() +
  geom_point(data = close_sub10, color = 'blue', size = 2) +
  geom_point(aes(106204113, 0.823), color = 'red', size = 3) +
  geom_vline(xintercept = c(106204113 + 15000, 106204113 - 15000), color = 'orange', linetype = 'dashed') +
  theme_pubr() +
  geom_rect(aes(xmin = 106202680, xmax = 106209408, ymin = -0.1, ymax = -0.05))

most_sub20 |>
  bind_rows(close_sub10) |>
  distinct(POS, .keep_all = TRUE) |>
  mutate(distance = POS - g2r_position) |>
  left_join(allele_sub) |>
  mutate(position = str_glue('chr14:{POS}')) |>
  write_csv('Archive/covid19/results/CDX_most_divergent_SNP.csv')

highest_CDX_fst |>
  left_join(chr14[1:5], by = join_by(position == POS)) |>
  select(-ID) |>
  mutate(distance = position - g2r_position,
         genome_coord = str_glue('chr14:{position}')) |>
  write_csv('Archive/covid19/results/CDX_highest_Fst_SNP.csv')
